# 环境准备
import torch
from transformers import BertTokenizer, BertForSequenceClassification

# 如果该文件改名copy.py就会报错！！！！！！
# 设备设置
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# 使用更优的中文BERT模型
MODEL_NAME = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)

# 模型定义
model = BertForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=2
).to(device)
torch.backends.cudnn.benchmark = True  # 启用cudnn优化


# 加载最佳模型
model.load_state_dict(torch.load('./model/best_model.bin'))
model = model.to(device)
print("模型加载完毕！")
# 预测函数
def predict(text, model, tokenizer, device, max_len=128):
    encoding = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        max_length=max_len,
        return_token_type_ids=False,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt',
    )

    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)

    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)

    logits = outputs.logits
    _, prediction = torch.max(logits, dim=1)
    
    return '攻击性言论' if prediction.item() == 1 else '正常言论'

# 测试示例
# test_texts = [
#     "傻子",
#     "成都的火锅文化确实很有特色",
#     "如果政策上恢复黑人昆仑奴的一切待遇,我想这黑人会被人间蒸发"
# ]

# for text in test_texts:
    

text = ""
while text != "q":
    print("请输入评论：")
    text = input()
    if text=="q":
        break
    print(f'"{text}" => {predict(text, model, tokenizer, device)}')